# ui/vertical_visualization.py
from PySide6.QtWidgets import (
    QDockWidget, QWidget, QVBoxLayout, QToolBar, QComboBox, QLabel,
    QSpinBox, QPushButton, QListWidget, QListWidgetItem, QColorDialog,
    QDoubleSpinBox, QSplitter, QMessageBox, QHBoxLayout, QMenu,
    QFileDialog, QInputDialog, QTabWidget, QSlider, QCheckBox, QLineEdit
)
from PySide6.QtCore import Qt, Signal, QPoint
from PySide6.QtGui import QColor, QIcon, QBrush, QAction
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import pandas as pd
import numpy as np
import os
from .base_widget import BaseWidgetMixin
from .curve_data_dialog import CurveDataDialog

# 设置matplotlib中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Zen Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


class VerticalVisualizationDock(QDockWidget):
    """竖向数据可视化组件"""

    dock_area = Qt.RightDockWidgetArea
    export_requested = Signal(str, pd.DataFrame)

    def __init__(self, title='竖向数据可视化', parent=None):
        # 直接调用QDockWidget的初始化方法
        super().__init__(title, parent)

        # 混入BaseWidgetMixin的功能
        if hasattr(BaseWidgetMixin, '__init__'):
            BaseWidgetMixin.__init__(self)

        # 设置允许浮动和关闭
        self.setFeatures(
            QDockWidget.DockWidgetClosable |
            QDockWidget.DockWidgetMovable |
            QDockWidget.DockWidgetFloatable
        )

        # 当前数据
        self.current_well_name = ""
        self.current_data = None
        self.depth_column = None
        self.selected_curve = None
        self.highlighted_range = None

        # 初始化UI
        self._init_ui()

    def _init_ui(self):
        """初始化UI"""
        main_widget = QWidget()
        main_layout = QVBoxLayout(main_widget)

        # 工具栏
        toolbar = QToolBar("竖向可视化工具")
        main_layout.addWidget(toolbar)

        # 图表类型选择
        toolbar.addWidget(QLabel("图表类型:"))
        self.plot_type_combo = QComboBox()
        self.plot_type_combo.addItems(["曲线图", "柱状图", "散点图", "阶梯图", "填充图"])
        self.plot_type_combo.currentTextChanged.connect(self._on_plot_type_changed)
        toolbar.addWidget(self.plot_type_combo)

        # 深度范围控制
        toolbar.addWidget(QLabel("深度范围:"))
        self.depth_min = QSpinBox()
        self.depth_min.setRange(0, 10000)
        self.depth_min.setSuffix(" m")
        self.depth_min.valueChanged.connect(self._update_plot)
        toolbar.addWidget(self.depth_min)

        toolbar.addWidget(QLabel("至"))

        self.depth_max = QSpinBox()
        self.depth_max.setRange(0, 10000)
        self.depth_max.setSuffix(" m")
        self.depth_max.valueChanged.connect(self._update_plot)
        toolbar.addWidget(self.depth_max)

        # 自动范围按钮
        auto_range_btn = QPushButton("自动范围")
        auto_range_btn.clicked.connect(self._auto_range)
        toolbar.addWidget(auto_range_btn)

        toolbar.addSeparator()

        # 交互模式按钮
        self.zoom_mode_btn = QAction("缩放", self)
        self.zoom_mode_btn.setCheckable(True)
        self.zoom_mode_btn.triggered.connect(lambda checked: self.canvas.set_interaction_mode('zoom', checked))
        toolbar.addAction(self.zoom_mode_btn)

        self.pan_mode_btn = QAction("平移", self)
        self.pan_mode_btn.setCheckable(True)
        self.pan_mode_btn.triggered.connect(lambda checked: self.canvas.set_interaction_mode('pan', checked))
        toolbar.addAction(self.pan_mode_btn)

        self.select_mode_btn = QAction("选择", self)
        self.select_mode_btn.setCheckable(True)
        self.select_mode_btn.triggered.connect(lambda checked: self.canvas.set_interaction_mode('select', checked))
        toolbar.addAction(self.select_mode_btn)

        toolbar.addSeparator()

        # 导出按钮
        export_btn = QAction("导出图像", self)
        export_btn.triggered.connect(self._export_image)
        toolbar.addAction(export_btn)

        export_data_btn = QAction("导出数据", self)
        export_data_btn.triggered.connect(self._export_data)
        toolbar.addAction(export_data_btn)

        # 创建分割器
        splitter = QSplitter(Qt.Vertical)

        # 图表画布
        self.canvas = VerticalPlotCanvas(parent=self)
        self.canvas.range_selected.connect(self._on_range_selected)
        splitter.addWidget(self.canvas)

        # 控制面板
        control_panel = QTabWidget()

        # 曲线选择面板
        curve_panel = QWidget()
        curve_layout = QVBoxLayout(curve_panel)

        curve_filter_layout = QHBoxLayout()
        curve_filter_layout.addWidget(QLabel("搜索:"))
        self.curve_filter = QLineEdit()
        self.curve_filter.textChanged.connect(self._filter_curves)
        curve_filter_layout.addWidget(self.curve_filter)
        curve_layout.addLayout(curve_filter_layout)

        self.curve_list = QListWidget()
        self.curve_list.setSelectionMode(QListWidget.ExtendedSelection)
        self.curve_list.itemChanged.connect(self._on_curve_visibility_changed)
        self.curve_list.itemClicked.connect(self._on_curve_selected)
        self.curve_list.setContextMenuPolicy(Qt.CustomContextMenu)
        self.curve_list.customContextMenuRequested.connect(self._show_curve_context_menu)
        curve_layout.addWidget(self.curve_list)

        control_panel.addTab(curve_panel, "曲线选择")

        # 样式设置面板
        style_panel = QWidget()
        style_layout = QVBoxLayout(style_panel)

        # 线宽设置
        line_width_layout = QHBoxLayout()
        line_width_layout.addWidget(QLabel("线宽:"))
        self.line_width_spin = QDoubleSpinBox()
        self.line_width_spin.setRange(0.1, 10.0)
        self.line_width_spin.setSingleStep(0.5)
        self.line_width_spin.setValue(1.0)
        self.line_width_spin.valueChanged.connect(self._on_line_width_changed)
        line_width_layout.addWidget(self.line_width_spin)
        style_layout.addLayout(line_width_layout)

        # 颜色设置
        color_layout = QHBoxLayout()
        color_layout.addWidget(QLabel("颜色:"))
        self.color_btn = QPushButton("选择颜色")
        self.color_btn.clicked.connect(self._on_color_changed)
        color_layout.addWidget(self.color_btn)
        style_layout.addLayout(color_layout)

        # 标记设置
        marker_layout = QHBoxLayout()
        marker_layout.addWidget(QLabel("标记:"))
        self.marker_combo = QComboBox()
        self.marker_combo.addItems(["None", "o", "s", "D", "^", "v", "*", "x", "+"])
        self.marker_combo.currentTextChanged.connect(self._on_marker_changed)
        marker_layout.addWidget(self.marker_combo)
        style_layout.addLayout(marker_layout)

        # 标记大小设置
        marker_size_layout = QHBoxLayout()
        marker_size_layout.addWidget(QLabel("标记大小:"))
        self.marker_size_spin = QDoubleSpinBox()
        self.marker_size_spin.setRange(1.0, 20.0)
        self.marker_size_spin.setSingleStep(1.0)
        self.marker_size_spin.setValue(5.0)
        self.marker_size_spin.valueChanged.connect(self._on_marker_size_changed)
        marker_size_layout.addWidget(self.marker_size_spin)
        style_layout.addLayout(marker_size_layout)

        # 透明度设置
        alpha_layout = QHBoxLayout()
        alpha_layout.addWidget(QLabel("透明度:"))
        self.alpha_slider = QSlider(Qt.Horizontal)
        self.alpha_slider.setRange(10, 100)
        self.alpha_slider.setValue(100)
        self.alpha_slider.valueChanged.connect(self._on_alpha_changed)
        alpha_layout.addWidget(self.alpha_slider)
        style_layout.addLayout(alpha_layout)

        # 对数坐标
        log_layout = QHBoxLayout()
        self.log_x_check = QCheckBox("X轴对数坐标")
        self.log_x_check.toggled.connect(self._on_log_scale_changed)
        log_layout.addWidget(self.log_x_check)
        style_layout.addLayout(log_layout)

        control_panel.addTab(style_panel, "样式设置")

        splitter.addWidget(control_panel)

        # 设置分割器比例
        splitter.setSizes([700, 300])

        main_layout.addWidget(splitter)
        self.setWidget(main_widget)

    def update_data(self, well_name, data):
        """更新数据"""
        self.current_well_name = well_name
        self.current_data = data
        self.highlighted_range = None

        if data is None or data.empty:
            self._clear_plot()
            return

        # 检测深度列
        self.depth_column = self._detect_depth_column(data.columns)
        if not self.depth_column:
            self.show_warning("未找到深度列（斜深/井深/TVD/MD）")
            return

        # 更新深度范围控件
        min_depth = int(data[self.depth_column].min())
        max_depth = int(data[self.depth_column].max())

        self.depth_min.setRange(min_depth, max_depth)
        self.depth_min.setValue(min_depth)

        self.depth_max.setRange(min_depth, max_depth)
        self.depth_max.setValue(max_depth)

        # 更新曲线列表
        self._update_curve_list()

        # 更新图表
        self._update_plot()

    def _clear_plot(self):
        """清除图表"""
        self.canvas.clear()
        self.curve_list.clear()

    def _update_plot(self):
        """更新图表"""
        if self.current_data is None or self.current_data.empty or not self.depth_column:
            return

        # 获取当前选择的深度范围
        depth_min = self.depth_min.value()
        depth_max = self.depth_max.value()

        # 确保深度范围有效
        if depth_min >= depth_max:
            self.show_warning("起始深度必须小于结束深度")
            return

        # 过滤数据
        filtered_data = self.current_data[
            (self.current_data[self.depth_column] >= depth_min) &
            (self.current_data[self.depth_column] <= depth_max)
        ]

        # 设置图表类型
        plot_type = self.plot_type_combo.currentText()
        self.canvas.set_plot_type(plot_type)

        # 获取可见曲线
        visible_curves = []
        for i in range(self.curve_list.count()):
            item = self.curve_list.item(i)
            if item.checkState() == Qt.Checked:
                visible_curves.append(item.text())

        # 清除当前图表
        self.canvas.clear()
        self.canvas.set_log_scale(self.log_x_check.isChecked())

        # 绘制所有可见曲线
        for curve_name in visible_curves:
            self.canvas.plot(
                y=filtered_data[self.depth_column],
                x=filtered_data[curve_name],
                label=curve_name
            )

        # 高亮显示选中区域
        if self.highlighted_range:
            depth_min, depth_max = self.highlighted_range
            self.canvas.highlight_range(depth_min, depth_max)

        # 重绘画布
        self.canvas.draw()

    def _update_curve_list(self):
        """更新曲线列表"""
        self.curve_list.clear()

        if self.current_data is None or self.depth_column is None:
            return

        # 添加所有曲线到列表
        for col in self.current_data.columns:
            if col != self.depth_column:
                item = QListWidgetItem(col)
                item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
                item.setCheckState(Qt.Checked)
                self.curve_list.addItem(item)

    def _auto_range(self):
        """自动设置深度范围"""
        if self.current_data is not None and self.depth_column is not None:
            min_depth = int(self.current_data[self.depth_column].min())
            max_depth = int(self.current_data[self.depth_column].max())

            self.depth_min.setValue(min_depth)
            self.depth_max.setValue(max_depth)

    def _filter_curves(self, text):
        """过滤曲线列表"""
        for i in range(self.curve_list.count()):
            item = self.curve_list.item(i)
            if text.lower() in item.text().lower():
                item.setHidden(False)
            else:
                item.setHidden(True)

    def _on_plot_type_changed(self, plot_type):
        """图表类型变化"""
        self.canvas.set_plot_type(plot_type)
        self._update_plot()

    def _on_curve_visibility_changed(self):
        """曲线可见性变化"""
        self._update_plot()

    def _on_curve_selected(self, item):
        """曲线选择变化"""
        self.selected_curve = item.text()

        # 更新样式控件
        if self.selected_curve and self.selected_curve in self.canvas.curves:
            style = self.canvas.curves[self.selected_curve]['style']
            self.line_width_spin.setValue(style['linewidth'])
            self.marker_combo.setCurrentText(style['marker'])
            self.marker_size_spin.setValue(style['markersize'])
            self.alpha_slider.setValue(int(style['alpha'] * 100))

    def _on_line_width_changed(self):
        """线宽变化"""
        self._apply_style_to_selected_curves(linewidth=self.line_width_spin.value())

    def _on_color_changed(self):
        """颜色变化"""
        color = QColorDialog.getColor()
        if color.isValid():
            color_hex = color.name()
            self._apply_style_to_selected_curves(color=color_hex)

    def _on_marker_changed(self):
        """标记变化"""
        self._apply_style_to_selected_curves(marker=self.marker_combo.currentText())

    def _on_marker_size_changed(self):
        """标记大小变化"""
        self._apply_style_to_selected_curves(markersize=self.marker_size_spin.value())

    def _on_alpha_changed(self):
        """透明度变化"""
        alpha = self.alpha_slider.value() / 100.0
        self._apply_style_to_selected_curves(alpha=alpha)

    def _on_log_scale_changed(self):
        """对数坐标变化"""
        self.canvas.set_log_scale(self.log_x_check.isChecked())
        self._update_plot()

    def _apply_style_to_selected_curves(self, **kwargs):
        """应用样式到选中的曲线"""
        selected_items = self.curve_list.selectedItems()
        if not selected_items:
            return

        for item in selected_items:
            curve_name = item.text()
            self.canvas.set_curve_style(curve_name, **kwargs)

        self._update_plot()

    def _detect_depth_column(self, columns):
        """检测深度列"""
        for col in ['斜深', '井深', 'TVD', 'MD', 'DEPTH']:
            if col in columns:
                return col
        return None

    def _show_curve_context_menu(self, position):
        """显示曲线上下文菜单"""
        if not self.curve_list.selectedItems():
            return

        menu = QMenu()

        # 添加复制选项
        copy_action = QAction("复制曲线", self)
        copy_action.triggered.connect(self._copy_curve)
        menu.addAction(copy_action)

        # 添加重命名选项
        rename_action = QAction("重命名曲线", self)
        rename_action.triggered.connect(self._rename_curve)
        menu.addAction(rename_action)

        # 添加删除选项
        delete_action = QAction("隐藏曲线", self)
        delete_action.triggered.connect(self._hide_curve)
        menu.addAction(delete_action)

        menu.addSeparator()

        # 添加查看数据选项
        view_data_action = QAction("查看详细数据", self)
        view_data_action.triggered.connect(self._view_curve_data)
        menu.addAction(view_data_action)

        menu.exec_(self.curve_list.mapToGlobal(position))

    def _copy_curve(self):
        """复制选中的曲线"""
        selected_items = self.curve_list.selectedItems()
        if not selected_items:
            return

        for item in selected_items:
            curve_name = item.text()
            new_name = f"{curve_name}_副本"

            # 确保新名称唯一
            counter = 1
            while new_name in [self.curve_list.item(i).text() for i in range(self.curve_list.count())]:
                new_name = f"{curve_name}_副本{counter}"
                counter += 1

            # 添加新曲线
            new_item = QListWidgetItem(new_name)
            new_item.setFlags(new_item.flags() | Qt.ItemIsUserCheckable)
            new_item.setCheckState(Qt.Checked)
            self.curve_list.addItem(new_item)

            # 如果原始曲线可见，则复制样式
            if item.checkState() == Qt.Checked:
                if curve_name in self.canvas.curves:
                    style = self.canvas.curves[curve_name]['style'].copy()
                    self.canvas.curves[new_name] = {'style': style, 'lines': []}

        self._update_plot()

    def _rename_curve(self):
        """重命名选中的曲线"""
        selected_items = self.curve_list.selectedItems()
        if not selected_items or len(selected_items) > 1:
            return

        item = selected_items[0]
        curve_name = item.text()

        new_name, ok = QInputDialog.getText(
            self, "重命名曲线", "请输入新的曲线名称:", text=curve_name
        )

        if ok and new_name and new_name != curve_name:
            # 确保新名称唯一
            if new_name in [self.curve_list.item(i).text() for i in range(self.curve_list.count())]:
                QMessageBox.warning(self, "警告", "已存在同名曲线!")
                return

            # 重命名曲线
            item.setText(new_name)

            # 更新图表中的曲线名称
            if curve_name in self.canvas.curves:
                self.canvas.curves[new_name] = self.canvas.curves.pop(curve_name)
                self.canvas.curves[new_name]['style']['label'] = new_name

            self._update_plot()

    def _hide_curve(self):
        """隐藏选中的曲线"""
        selected_items = self.curve_list.selectedItems()
        if not selected_items:
            return

        for item in selected_items:
            item.setCheckState(Qt.Unchecked)

        self._update_plot()

    def _view_curve_data(self):
        """查看曲线详细数据"""
        if not self.current_data or not self.depth_column:
            return

        selected_items = self.curve_list.selectedItems()
        if not selected_items:
            return

        # 获取当前显示的数据
        depth_min = self.depth_min.value()
        depth_max = self.depth_max.value()

        filtered_data = self.current_data[
            (self.current_data[self.depth_column] >= depth_min) &
            (self.current_data[self.depth_column] <= depth_max)
        ]

        if filtered_data.empty:
            QMessageBox.warning(self, "警告", "没有数据可供显示")
            return

        # 如果有选中区域，只显示选中区域的数据
        if self.highlighted_range:
            depth_min, depth_max = self.highlighted_range
            filtered_data = filtered_data[
                (filtered_data[self.depth_column] >= depth_min) &
                (filtered_data[self.depth_column] <= depth_max)
            ]

        # 创建数据对话框
        curve_names = [item.text() for item in selected_items]
        columns_to_show = [self.depth_column] + curve_names
        data_to_show = filtered_data[columns_to_show]

        dialog = CurveDataDialog(
            data_to_show,
            self,
            f"{self.current_well_name} - {', '.join(curve_names)} 详细数据"
        )

        # 连接导出信号
        dialog.data_exported.connect(lambda path: self.export_requested.emit(
            ', '.join(curve_names), data_to_show))

        dialog.exec_()

    def _export_image(self):
        """导出图像"""
        file_dialog = QFileDialog()
        file_path, _ = file_dialog.getSaveFileName(
            self, "导出图像", "", "PNG Images (*.png);;JPEG Images (*.jpg);;SVG Images (*.svg)"
        )
        if file_path:
            self.canvas.figure.savefig(file_path)

    def _export_data(self):
        """导出数据"""
        if self.current_data is None or self.current_data.empty:
            QMessageBox.warning(self, "警告", "没有数据可供导出")
            return

        file_dialog = QFileDialog()
        file_path, _ = file_dialog.getSaveFileName(
            self, "导出数据", "", "CSV Files (*.csv);;Excel Files (*.xlsx)"
        )
        if file_path:
            if file_path.endswith('.csv'):
                self.current_data.to_csv(file_path, index=False)
            elif file_path.endswith('.xlsx'):
                self.current_data.to_excel(file_path, index=False)

    def _on_range_selected(self, depth_min, depth_max):
        """范围选择变化"""
        self.highlighted_range = (depth_min, depth_max)
        self._update_plot()

    def show_warning(self, message):
        """显示警告消息框"""
        QMessageBox.warning(self, "警告", message)


class VerticalPlotCanvas(FigureCanvas):
    """竖向绘图画布"""
    range_selected = Signal(float, float)

    def __init__(self, parent=None):
        fig = Figure()
        self.axes = fig.add_subplot(111)
        super().__init__(fig)
        self.setParent(parent)

        self.curves = {}
        self.plot_type = "曲线图"
        self.log_scale = False
        self.interaction_mode = None

        self._setup_interaction()

    def _setup_interaction(self):
        """设置交互模式"""
        self.mpl_connect('button_press_event', self._on_mouse_press)
        self.mpl_connect('button_release_event', self._on_mouse_release)
        self.mpl_connect('motion_notify_event', self._on_mouse_move)

    def set_plot_type(self, plot_type):
        """设置图表类型"""
        self.plot_type = plot_type

    def set_log_scale(self, log_scale):
        """设置对数坐标"""
        self.log_scale = log_scale
        if log_scale:
            self.axes.set_xscale('log')
        else:
            self.axes.set_xscale('linear')

    def set_interaction_mode(self, mode, enabled):
        """设置交互模式"""
        if enabled:
            if self.interaction_mode:
                getattr(self, f'_disable_{self.interaction_mode}_mode')()
            self.interaction_mode = mode
            getattr(self, f'_enable_{self.interaction_mode}_mode')()
        else:
            if self.interaction_mode == mode:
                getattr(self, f'_disable_{self.interaction_mode}_mode')()
                self.interaction_mode = None

    def _enable_zoom_mode(self):
        """启用缩放模式"""
        self.fig.canvas.toolbar.zoom()

    def _disable_zoom_mode(self):
        """禁用缩放模式"""
        self.fig.canvas.toolbar.zoom()

    def _enable_pan_mode(self):
        """启用平移模式"""
        self.fig.canvas.toolbar.pan()

    def _disable_pan_mode(self):
        """禁用平移模式"""
        self.fig.canvas.toolbar.pan()

    def _enable_select_mode(self):
        """启用选择模式"""
        self.select_start = None
        self.select_end = None
        self.select_rect = None

    def _disable_select_mode(self):
        """禁用选择模式"""
        if self.select_rect:
            self.select_rect.remove()
            self.select_rect = None
        self.select_start = None
        self.select_end = None

    def _on_mouse_press(self, event):
        """鼠标按下事件"""
        if self.interaction_mode == 'select':
            if event.inaxes == self.axes:
                self.select_start = event.ydata

    def _on_mouse_release(self, event):
        """鼠标释放事件"""
        if self.interaction_mode == 'select':
            if event.inaxes == self.axes and self.select_start is not None:
                self.select_end = event.ydata
                depth_min = min(self.select_start, self.select_end)
                depth_max = max(self.select_start, self.select_end)
                self.range_selected.emit(depth_min, depth_max)
                self._disable_select_mode()

    def _on_mouse_move(self, event):
        """鼠标移动事件"""
        if self.interaction_mode == 'select':
            if event.inaxes == self.axes and self.select_start is not None:
                if self.select_rect:
                    self.select_rect.remove()
                y1 = min(self.select_start, event.ydata)
                y2 = max(self.select_start, event.ydata)
                self.select_rect = self.axes.axhspan(y1, y2, facecolor='gray', alpha=0.3)
                self.draw()

    def clear(self):
        """清除图表"""
        self.axes.clear()
        self.curves = {}
        self.set_log_scale(self.log_scale)

    def plot(self, x, y, label):
        """绘制曲线"""
        if self.plot_type == "曲线图":
            line, = self.axes.plot(x, y, label=label)
        elif self.plot_type == "柱状图":
            line = self.axes.bar(x, y, label=label)
        elif self.plot_type == "散点图":
            line, = self.axes.scatter(x, y, label=label)
        elif self.plot_type == "阶梯图":
            line, = self.axes.step(x, y, label=label)
        elif self.plot_type == "填充图":
            line = self.axes.fill_betweenx(y, x, label=label)

        style = {
            'linewidth': 1.0,
            'color': 'blue',
            'marker': 'None',
            'markersize': 5.0,
            'alpha': 1.0
        }
        self.curves[label] = {'style': style, 'lines': [line]}
        self.axes.legend()

    def set_curve_style(self, curve_name, **kwargs):
        """设置曲线样式"""
        if curve_name in self.curves:
            style = self.curves[curve_name]['style']
            style.update(kwargs)
            lines = self.curves[curve_name]['lines']
            for line in lines:
                if isinstance(line, plt.Line2D):
                    line.set_linewidth(style['linewidth'])
                    line.set_color(style['color'])
                    line.set_marker(style['marker'])
                    line.set_markersize(style['markersize'])
                    line.set_alpha(style['alpha'])
                elif isinstance(line, plt.Polygon):
                    line.set_facecolor(style['color'])
                    line.set_alpha(style['alpha'])

    def highlighted_range(self, depth_min, depth_max):
        """高亮显示选中区域"""
        self.axes.axhspan(depth_min, depth_max, facecolor='gray', alpha=0.3)
